import numpy as np
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# 1.	加载西瓜数据(20分)
# load
data = np.loadtxt('xigua.txt', delimiter=',')
m = len(data)
# shuffle
np.random.seed(1)
np.random.shuffle(data)
# preprocessing
x = data[:, :-1]
y = data[:, -1]
# scale
x = StandardScaler().fit_transform(x)
# split
x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7)

# 2.	调用决策树分类模型，并训练模型(40分)
clf2 = DecisionTreeClassifier(max_depth=2)
clf2.fit(x_train, y_train)
clf5 = DecisionTreeClassifier(max_depth=5)
clf5.fit(x_train, y_train)

# 3.	计算出深度=2和5时的准确率(40分)
print(f'深度=2时的准确率: 训练时为{clf2.score(x_train, y_train)}, 测试时为{clf2.score(x_test, y_test)}')
print(f'深度=5时的准确率: 训练时为{clf5.score(x_train, y_train)}, 测试时为{clf5.score(x_test, y_test)}')
